import os
import xml.etree.ElementTree as ET
import matplotlib.pyplot as plt
import torchvision.transforms as T
import numpy as np
from pathlib import Path
from PIL import Image
from data_check import check_annotations
import warnings


def show(imgs):
    fix, axs = plt.subplots(ncols=len(imgs), squeeze=False)
    for i, img in enumerate(imgs):
        img = T.ToPILImage()(img.to('cpu'))
        axs[0, i].imshow(np.asarray(img))
        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])


def plot(imgs, with_orig=False, row_title=None, **imshow_kwargs):
    if not isinstance(imgs[0], list):
        # Make a 2d grid even if there's just 1 row
        imgs = [imgs]

    num_rows = len(imgs)
    # num_cols = len(imgs[0]) + with_orig
    num_cols = len(imgs[0]) + with_orig
    fig, axs = plt.subplots(nrows=num_rows, ncols=num_cols, squeeze=False)
    for row_idx, row in enumerate(imgs):
        row = [orig_img] + row if with_orig else row
        for col_idx, img in enumerate(row):
            ax = axs[row_idx, col_idx]
            ax.imshow(np.asarray(img), **imshow_kwargs)
            ax.set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])

    if with_orig:
        axs[0, 0].set(title='Original image')
        axs[0, 0].title.set_size(8)
    if row_title is not None:
        for row_idx in range(num_rows):
            axs[row_idx, 0].set(ylabel=row_title[row_idx])
    plt.tight_layout()


def save_changed_jpg_xml(imgs, change_method=None, filename=None):
    new_name_list = []
    for i in range(len(imgs)):
        width, height = imgs[i].size
        fig = plt.figure(figsize=(width / 100.0, height / 100.0), dpi=130)
        plt.axis('off')
        plt.imshow(np.asarray(imgs[i]))
        new_name = change_method + filename.split(".")[0] + "_" + str(i) + ".jpg"
        new_name_list.append(new_name)
        plt.savefig(os.path.join("MaskDatasets/augmentation_jpgs", new_name), pad_inches=0.0)
        plt.close(fig)

    xml_path = os.path.join("MaskDatasets", "Annotations")
    augmentation_jpgs_path = os.path.join("MaskDatasets", "augmentation_jpgs")
    augmentation_xmls_path = os.path.join("MaskDatasets", "augmentation_xmls")

    aug_jpgs = new_name_list
    new_xml_name = []
    for i in range(len(aug_jpgs)):
        new_xml_name.append(aug_jpgs[i].split('.')[0] + ".xml")
    orig_xml = os.path.join(xml_path, filename.split('.')[0] + ".xml")

    for i in range(len(aug_jpgs)):
        os.popen('copy %s %s' % (orig_xml, os.path.join(augmentation_xmls_path, new_xml_name[i])))


def save_changed_xml(filename):
    # poor_mask = check_annotations("MaskDatasets")[1]
    xml_path = os.path.join("MaskDatasets", "Annotations")
    augmentation_jpgs_path = os.path.join("MaskDatasets", "augmentation_jpgs")
    augmentation_xmls_path = os.path.join("MaskDatasets", "augmentation_xmls")

    aug_jpgs = os.listdir(      )
    new_xml_name = []
    for i in range(len(aug_jpgs)):
        new_xml_name.append(aug_jpgs[i].split('.')[0] + ".xml")
    orig_xml = os.path.join(xml_path, filename.split('.')[0] + ".xml")

    for i in range(len(aug_jpgs)):
        os.popen('copy %s %s' % (orig_xml, os.path.join(augmentation_xmls_path, new_xml_name[i])))


def xmls2jpgxmls():
    xml_path = os.path.join("MaskDatasets", "augmentation_xmls")
    xmls = os.listdir(xml_path)
    print(xmls)
    # 打开xml文档
    for xml in xmls:
        doc = ET.parse(os.path.join(xml_path, xml))
        root = doc.getroot()
        sub = root.find('filename')  # 找到filename标签，
        sub.text = xml.split(".")[0] + ".jpg"  # 修改标签内容
        doc.write(os.path.join(xml_path, xml))  # 保存修改


if __name__ == '__main__':
    warnings.filterwarnings("ignore")
    plt.rcParams["savefig.bbox"] = 'tight'
    poor_mask = check_annotations("MaskDatasets")[1]
    for file in poor_mask:
        file = file + ".jpg"
        orig_img = Image.open(Path('MaskDatasets/JPEGImages') / file)

        # Resize : 改变图像大小 标记框位置改变
        # resized_imgs = [T.Resize(size=size)(orig_img) for size in (30, 50, 100)]
        # save_changed_jpg(resized_imgs, change_method="Resize", filename=file)
        # plot(resized_imgs)

        # Grayscale : 转灰度图像 标记框位置不变
        # gray_img = T.Grayscale()(orig_img)
        # plot([gray_img], cmap='gray')

        # ColorJitter : 随机改变图像的亮度、饱和度和其他属性 标记框位置不变
        print(orig_img.size)
        jitter = T.ColorJitter(brightness=.5, hue=.5, contrast=.5)
        jitted_imgs = [jitter(orig_img) for _ in range(12)]
        save_changed_jpg_xml(jitted_imgs, change_method="ColorJitter", filename=file)
        # plot(jitted_imgs)

        # GaussianBlur : 高斯模糊变换图像 标记框位置不变
        blurrer = T.GaussianBlur(kernel_size=(3, 9), sigma=(0.1, 10))
        blurred_imgs = [blurrer(orig_img) for _ in range(4)]
        save_changed_jpg_xml(blurred_imgs, change_method="GaussianBlur", filename=file)
        # plot(blurred_imgs)

        # RandomAdjustSharpness : 随机调整给定图像的锐度 标记框位置不变
        sharpness_adjuster = T.RandomAdjustSharpness(sharpness_factor=2)
        sharpened_imgs = [sharpness_adjuster(orig_img) for _ in range(4)]
        save_changed_jpg_xml(sharpened_imgs, change_method="RandomAdjustSharpness", filename=file)
        # plot(sharpened_imgs)

        # RandomSolarize : 通过反转阈值以上的所有像素值来随机曝光图像 标记框位置不变
        solarizer = T.RandomSolarize(threshold=192.0)
        solarized_imgs = [solarizer(orig_img) for _ in range(4)]
        save_changed_jpg_xml(solarized_imgs, change_method="RandomSolarize", filename=file)
        # plot(solarized_imgs)

        # RandomPosterize : 通过减少每个颜色通道的位数来随机对图像进行分色 标记框位置不变
        posterizer = T.RandomPosterize(bits=2)
        posterized_imgs = [posterizer(orig_img) for _ in range(6)]
        save_changed_jpg_xml(posterized_imgs, change_method="RandomPosterize", filename=file)
        # plot(posterized_imgs)
        # plt.show()
        # save_changed_xml(filename=file)
    xmls2jpgxmls()
